# pip install pillow opencv-python numpy

import numpy as np
import random
from collections import deque
import argparse
import os
import datetime
from PIL import Image, ImageDraw, ImageFont

# 颜色定义 (RGB格式)
BACKGROUND = (240, 240, 245)
WALL_COLOR = (30, 30, 60)
PATH_COLOR = (100, 200, 255)
SOLUTION_COLOR = (255, 100, 100)  # 红色
START_COLOR = (50, 200, 50)      # 绿色
END_COLOR = (255, 50, 50)        # 红色
TEXT_COLOR = (30, 30, 60)

class MazeGenerator:
    def __init__(self, rows=10, cols=10):
        self.rows = rows
        self.cols = cols
        self.maze = []
        self.solution_path = []
        self.visited = []
        self.start = (0, 0)
        self.end = (rows-1, cols-1)
        self.generate_maze()
    
    def generate_maze(self):
        # 初始化迷宫网格
        self.maze = [[{'top': True, 'right': True, 'bottom': True, 'left': True} 
                     for _ in range(self.cols)] for _ in range(self.rows)]
        
        # 初始化访问数组
        self.visited = [[False for _ in range(self.cols)] for _ in range(self.rows)]
        
        # 从起点开始生成迷宫
        stack = [(self.start[0], self.start[1])]
        self.visited[self.start[0]][self.start[1]] = True
        
        # 方向：上、右、下、左
        directions = [(-1, 0), (0, 1), (1, 0), (0, -1)]
        dir_names = ['top', 'right', 'bottom', 'left']
        
        while stack:
            r, c = stack[-1]
            neighbors = []
            
            # 检查未访问的邻居
            for i, (dr, dc) in enumerate(directions):
                nr, nc = r + dr, c + dc
                if 0 <= nr < self.rows and 0 <= nc < self.cols and not self.visited[nr][nc]:
                    neighbors.append((i, nr, nc))
            
            if neighbors:
                # 随机选择一个邻居
                direction, nr, nc = random.choice(neighbors)
                
                # 打通墙壁
                if direction == 0:   # 上
                    self.maze[r][c]['top'] = False
                    self.maze[nr][nc]['bottom'] = False
                elif direction == 1: # 右
                    self.maze[r][c]['right'] = False
                    self.maze[nr][nc]['left'] = False
                elif direction == 2: # 下
                    self.maze[r][c]['bottom'] = False
                    self.maze[nr][nc]['top'] = False
                elif direction == 3: # 左
                    self.maze[r][c]['left'] = False
                    self.maze[nr][nc]['right'] = False
                
                # 标记为已访问并加入栈
                self.visited[nr][nc] = True
                stack.append((nr, nc))
            else:
                # 回溯
                stack.pop()
        
        # 找到解决方案
        self.find_solution()
    
    def find_solution(self):
        # 使用BFS找到从起点到终点的路径
        queue = deque([(self.start[0], self.start[1], [])])
        visited = [[False for _ in range(self.cols)] for _ in range(self.rows)]
        visited[self.start[0]][self.start[1]] = True
        
        while queue:
            r, c, path = queue.popleft()
            new_path = path + [(r, c)]
            
            if (r, c) == self.end:
                self.solution_path = new_path
                return
            
            # 检查可以移动的方向
            if not self.maze[r][c]['top'] and not visited[r-1][c]:
                visited[r-1][c] = True
                queue.append((r-1, c, new_path))
            
            if not self.maze[r][c]['right'] and not visited[r][c+1]:
                visited[r][c+1] = True
                queue.append((r, c+1, new_path))
            
            if not self.maze[r][c]['bottom'] and not visited[r+1][c]:
                visited[r+1][c] = True
                queue.append((r+1, c, new_path))
            
            if not self.maze[r][c]['left'] and not visited[r][c-1]:
                visited[r][c-1] = True
                queue.append((r, c-1, new_path))

def create_maze_image(maze_generator, cell_size=40, wall_thickness=2, show_solution=True):
    rows = maze_generator.rows
    cols = maze_generator.cols
    
    # 创建空白图像 (宽度, 高度)
    img_width = cols * cell_size + 100
    img_height = rows * cell_size + 150
    img = Image.new('RGB', (img_width, img_height), BACKGROUND)
    draw = ImageDraw.Draw(img)
    
    try:
        # 尝试加载字体
        font = ImageFont.truetype("arial.ttf", 24)
        small_font = ImageFont.truetype("arial.ttf", 18)
    except:
        # 如果找不到字体，使用默认字体
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()
    
    # 绘制标题
    title = f"Maze Generator ({rows}x{cols})"
    draw.text((img_width//2 - 100, 20), title, fill=TEXT_COLOR, font=font)
    
    # 绘制说明
    info = "Start(Green) -> End(Red), Unique Path"
    draw.text((img_width//2 - 120, 60), info, fill=TEXT_COLOR, font=small_font)
    
    # 绘制生成时间
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    time_text = f"Generated at: {timestamp}"
    draw.text((20, img_height - 30), time_text, fill=TEXT_COLOR, font=small_font)
    
    # 计算迷宫偏移量使其居中
    offset_x = (img_width - cols * cell_size) // 2
    offset_y = 100
    
    # 绘制迷宫
    for r in range(rows):
        for c in range(cols):
            x = offset_x + c * cell_size
            y = offset_y + r * cell_size
            
            # 绘制墙壁
            if maze_generator.maze[r][c]['top']:
                draw.line([(x, y), (x + cell_size, y)], fill=WALL_COLOR, width=wall_thickness)
            if maze_generator.maze[r][c]['right']:
                draw.line([(x + cell_size, y), (x + cell_size, y + cell_size)], fill=WALL_COLOR, width=wall_thickness)
            if maze_generator.maze[r][c]['bottom']:
                draw.line([(x, y + cell_size), (x + cell_size, y + cell_size)], fill=WALL_COLOR, width=wall_thickness)
            if maze_generator.maze[r][c]['left']:
                draw.line([(x, y), (x, y + cell_size)], fill=WALL_COLOR, width=wall_thickness)
    
    # 绘制解决方案路径
    if show_solution:
        for i in range(len(maze_generator.solution_path) - 1):
            r1, c1 = maze_generator.solution_path[i]
            r2, c2 = maze_generator.solution_path[i+1]
            
            x1 = offset_x + c1 * cell_size + cell_size // 2
            y1 = offset_y + r1 * cell_size + cell_size // 2
            x2 = offset_x + c2 * cell_size + cell_size // 2
            y2 = offset_y + r2 * cell_size + cell_size // 2
            
            draw.line([(x1, y1), (x2, y2)], fill=SOLUTION_COLOR, width=4)
    
    # 绘制起点和终点
    start_x = offset_x + maze_generator.start[1] * cell_size + cell_size // 2
    start_y = offset_y + maze_generator.start[0] * cell_size + cell_size // 2
    draw.ellipse([(start_x - cell_size//3, start_y - cell_size//3), 
                 (start_x + cell_size//3, start_y + cell_size//3)], fill=START_COLOR)
    
    end_x = offset_x + maze_generator.end[1] * cell_size + cell_size // 2
    end_y = offset_y + maze_generator.end[0] * cell_size + cell_size // 2
    draw.ellipse([(end_x - cell_size//3, end_y - cell_size//3), 
                 (end_x + cell_size//3, end_y + cell_size//3)], fill=END_COLOR)
    
    return img

def save_maze_image(maze_generator, output_dir, solution=True, cell_size=40, prefix="maze"):
    """生成并保存迷宫图片为PNG格式"""
    # 创建保存目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 生成文件名（包含时间戳和尺寸）
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    solution_suffix = "_solution" if solution else ""
    size = maze_generator.rows
    filename = f"{prefix}_{size}x{size}{solution_suffix}_{timestamp}.png"
    filepath = os.path.join(output_dir, filename)
    
    # 创建迷宫图像
    img = create_maze_image(maze_generator, cell_size=cell_size, show_solution=solution)
    
    # 保存图片
    img.save(filepath, "PNG")
    print(f"迷宫图片已保存至: {filepath}")
    
    return filepath

def generate_mazes(output_dir, size=10, count=1, solution=True, cell_size=40):
    """生成指定数量的迷宫并保存"""
    for i in range(count):
        print(f"正在生成迷宫 {i+1}/{count} ({size}x{size})...")
        generator = MazeGenerator(size, size)
        save_maze_image(generator, output_dir, solution, cell_size)

def main():
    parser = argparse.ArgumentParser(description='命令行迷宫生成器')
    parser.add_argument('--size', type=int, default=10, 
                        help='迷宫尺寸 (默认: 10)')
    parser.add_argument('--count', type=int, default=1, 
                        help='生成迷宫的数量 (默认: 1)')
    parser.add_argument('--output', type=str, default="maze_images", 
                        help='输出目录 (默认: maze_images)')
    parser.add_argument('--no-solution', action='store_true', 
                        help='不在图像中包含解决方案路径')
    parser.add_argument('--cell-size', type=int, default=40, 
                        help='每个单元格的像素大小 (默认: 40)')
    
    args = parser.parse_args()
    
    # 生成迷宫
    generate_mazes(
        output_dir=args.output,
        size=args.size,
        count=args.count,
        solution=not args.no_solution,
        cell_size=args.cell_size
    )
    
    print(f"成功生成 {args.count} 个 {args.size}x{args.size} 的迷宫，并保存到 '{args.output}' 目录")

if __name__ == "__main__":
    main()
